package org.study.clickhouse.api.util;

import com.clickhouse.client.*;
import com.clickhouse.data.*;
import com.clickhouse.data.format.BinaryStreamUtils;

import java.io.IOException;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

/**
 * 参考 https://github.com/ClickHouse/clickhouse-java/blob/main/examples/client/src/main/java/com/clickhouse/examples/jdbc/Main.java
 *
 * @author doveylovey
 * @version v1.0.0
 * @email 1135782208@qq.com
 * @date 2023年05月10日
 */
public class ClickHouseMainTest3 {
    public static void main(String[] args) throws ExecutionException, InterruptedException, IOException, ClickHouseException {
        ClickHouseNode server = ClickHouseNode.builder()
                .host("127.0.0.1")
                .port(ClickHouseProtocol.HTTP, 8123)
                // .port(ClickHouseProtocol.GRPC, 9000)
                // .port(ClickHouseProtocol.TCP, 9100)
                .database("default")
                .credentials(ClickHouseCredentials.fromUserAndPassword("default", "root"))
                .build();

        String table = "ck_table";
        dropAndCreateTable(server, table);
        System.out.println("Insert: " + insert(server, table));
        System.out.println("Query: " + query(server, table));
    }

    static void dropAndCreateTable(ClickHouseNode server, String table) throws ExecutionException, InterruptedException, ClickHouseException {
        try (ClickHouseClient client = ClickHouseClient.newInstance(server.getProtocol())) {
            ClickHouseRequest<?> request = client.read(server);
            // or use future chaining
            request.query("drop table if exists " + table)
                    .execute()
                    .get();
            request.query("create table " + table + "(a String, b Nullable(String)) engine=MergeTree() order by a")
                    .executeAndWait();
        }
    }

    static long insert(ClickHouseNode server, String table) throws IOException, ExecutionException, InterruptedException {
        try (ClickHouseClient client = ClickHouseClient.newInstance(server.getProtocol())) {
            ClickHouseRequest.Mutation request = client.read(server)
                    .write()
                    .table(table)
                    .format(ClickHouseFormat.RowBinary);
            ClickHouseConfig config = request.getConfig();
            CompletableFuture<ClickHouseResponse> future;
            // back-pressuring is not supported, you can adjust the first two arguments
            try (ClickHousePipedOutputStream stream = ClickHouseDataStreamFactory.getInstance()
                    .createPipedOutputStream(config, (Runnable) null)) {
                // in async mode, which is default, execution happens in a worker thread
                future = request.data(stream.getInputStream())
                        .execute();
                // writing happens in main thread
                for (int i = 0; i < 10_000; i++) {
                    BinaryStreamUtils.writeString(stream, String.valueOf(i % 16));
                    BinaryStreamUtils.writeNonNull(stream);
                    BinaryStreamUtils.writeString(stream, UUID.randomUUID().toString());
                }
            }
            // response should be always closed
            try (ClickHouseResponse response = future.get()) {
                ClickHouseResponseSummary summary = response.getSummary();
                return summary.getWrittenRows();
            }
        }
    }

    static int query(ClickHouseNode server, String table) throws ExecutionException, InterruptedException {
        try (ClickHouseClient client = ClickHouseClient.newInstance(server.getProtocol());
             ClickHouseResponse response = client.read(server)
                     // prefer to use RowBinaryWithNamesAndTypes as it's fully supported
                     // see details at https://github.com/ClickHouse/clickhouse-java/issues/928
                     .format(ClickHouseFormat.RowBinaryWithNamesAndTypes)
                     .query("select * from " + table)
                     .execute()
                     .get()
        ) {
            int count = 0;
            // or use stream API via response.stream()
            for (ClickHouseRecord record : response.records()) {
                count++;
            }
            return count;
        }
    }
}
